"""
视频帧提取器
实现基于OpenCV的RTSP/HTTP-FLV流解码功能
"""
import asyncio
import logging
import time
from typing import AsyncGenerator, Optional, Dict, Any
import cv2
import numpy as np
from contextlib import asynccontextmanager

from core.config import get_settings

logger = logging.getLogger(__name__)
settings = get_settings()


class VideoFrameExtractor:
    """视频帧提取器"""
    
    def __init__(self, stream_url: str, target_fps: int = None, camera_id: str = None):
        """
        初始化视频帧提取器
        
        Args:
            stream_url: 视频流URL
            target_fps: 目标帧率，默认使用配置中的值
            camera_id: 摄像头ID，用于日志记录
        """
        self.stream_url = stream_url
        self.target_fps = target_fps or settings.VIDEO_PROCESSING_FPS
        self.camera_id = camera_id or "unknown"
        self.cap: Optional[cv2.VideoCapture] = None
        self.is_running = False
        self.last_frame_time = 0
        self.frame_interval = 1.0 / self.target_fps
        self.reconnect_attempts = 0
        self.max_reconnect_attempts = settings.VIDEO_RECONNECT_ATTEMPTS
        self.reconnect_delay = settings.VIDEO_RECONNECT_DELAY
        
        # 性能统计
        self.stats = {
            "frames_extracted": 0,
            "frames_dropped": 0,
            "reconnections": 0,
            "last_successful_frame": None
        }
        
    async def __aenter__(self):
        """异步上下文管理器入口"""
        await self.start()
        return self
        
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """异步上下文管理器出口"""
        await self.stop()
        
    async def start(self) -> bool:
        """启动视频流捕获"""
        try:
            logger.info(f"启动摄像头 {self.camera_id} 的视频流捕获: {self.stream_url}")
            
            # 设置OpenCV捕获参数
            self.cap = cv2.VideoCapture(self.stream_url)
            
            # 配置捕获参数
            self._configure_capture()
            
            # 测试连接
            if not self.cap.isOpened():
                logger.error(f"无法打开摄像头 {self.camera_id} 的视频流: {self.stream_url}")
                return False
                
            # 尝试读取一帧来验证流
            ret, frame = self.cap.read()
            if not ret or frame is None:
                logger.error(f"无法从摄像头 {self.camera_id} 读取视频帧")
                return False
                
            self.is_running = True
            self.reconnect_attempts = 0
            logger.info(f"摄像头 {self.camera_id} 视频流启动成功")
            return True
            
        except Exception as e:
            logger.error(f"启动摄像头 {self.camera_id} 视频流失败: {str(e)}")
            return False
            
    async def stop(self):
        """停止视频流捕获"""
        try:
            self.is_running = False
            if self.cap:
                self.cap.release()
                self.cap = None
            logger.info(f"摄像头 {self.camera_id} 视频流已停止")
        except Exception as e:
            logger.error(f"停止摄像头 {self.camera_id} 视频流失败: {str(e)}")
            
    def _configure_capture(self):
        """配置视频捕获参数"""
        if not self.cap:
            return
            
        try:
            # 设置缓冲区大小
            self.cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
            
            # 设置超时
            self.cap.set(cv2.CAP_PROP_OPEN_TIMEOUT_MSEC, settings.FRAME_EXTRACTION_TIMEOUT * 1000)
            self.cap.set(cv2.CAP_PROP_READ_TIMEOUT_MSEC, settings.FRAME_EXTRACTION_TIMEOUT * 1000)
            
            # 尝试设置帧率
            self.cap.set(cv2.CAP_PROP_FPS, self.target_fps)
            
            # 获取实际的视频参数
            actual_fps = self.cap.get(cv2.CAP_PROP_FPS)
            width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
            height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
            
            logger.info(f"摄像头 {self.camera_id} 视频参数: {width}x{height}@{actual_fps}fps")
            
        except Exception as e:
            logger.warning(f"配置摄像头 {self.camera_id} 捕获参数失败: {str(e)}")
            
    async def extract_frames(self) -> AsyncGenerator[np.ndarray, None]:
        """
        异步提取视频帧
        
        Yields:
            numpy.ndarray: 视频帧数据
        """
        if not self.is_running or not self.cap:
            logger.error(f"摄像头 {self.camera_id} 视频流未启动")
            return
            
        while self.is_running:
            try:
                current_time = time.time()
                
                # 控制帧率
                if current_time - self.last_frame_time < self.frame_interval:
                    await asyncio.sleep(0.01)  # 短暂休眠
                    continue
                    
                # 读取帧
                ret, frame = self.cap.read()
                
                if not ret or frame is None:
                    logger.warning(f"摄像头 {self.camera_id} 读取帧失败，尝试重连")
                    
                    # 尝试重连
                    if await self._handle_stream_interruption():
                        continue
                    else:
                        break
                        
                # 更新统计信息
                self.stats["frames_extracted"] += 1
                self.stats["last_successful_frame"] = current_time
                self.last_frame_time = current_time
                
                yield frame
                
            except Exception as e:
                logger.error(f"摄像头 {self.camera_id} 提取帧时发生错误: {str(e)}")
                
                # 尝试重连
                if not await self._handle_stream_interruption():
                    break
                    
        logger.info(f"摄像头 {self.camera_id} 帧提取结束")
        
    async def _handle_stream_interruption(self) -> bool:
        """
        处理流中断，实现自动重连
        
        Returns:
            bool: 重连是否成功
        """
        if self.reconnect_attempts >= self.max_reconnect_attempts:
            logger.error(f"摄像头 {self.camera_id} 重连次数超过限制 ({self.max_reconnect_attempts})")
            return False
            
        self.reconnect_attempts += 1
        self.stats["reconnections"] += 1
        
        logger.info(f"摄像头 {self.camera_id} 开始第 {self.reconnect_attempts} 次重连")
        
        try:
            # 释放当前连接
            if self.cap:
                self.cap.release()
                
            # 等待重连延迟
            await asyncio.sleep(self.reconnect_delay)
            
            # 重新建立连接
            self.cap = cv2.VideoCapture(self.stream_url)
            self._configure_capture()
            
            # 测试连接
            if self.cap.isOpened():
                ret, frame = self.cap.read()
                if ret and frame is not None:
                    logger.info(f"摄像头 {self.camera_id} 重连成功")
                    self.reconnect_attempts = 0  # 重置重连计数
                    return True
                    
            logger.warning(f"摄像头 {self.camera_id} 重连失败")
            return False
            
        except Exception as e:
            logger.error(f"摄像头 {self.camera_id} 重连过程中发生错误: {str(e)}")
            return False
            
    def get_stream_info(self) -> Dict[str, Any]:
        """
        获取视频流信息
        
        Returns:
            Dict[str, Any]: 流信息
        """
        if not self.cap or not self.cap.isOpened():
            return {
                "status": "disconnected",
                "camera_id": self.camera_id,
                "stream_url": self.stream_url
            }
            
        try:
            return {
                "status": "connected" if self.is_running else "stopped",
                "camera_id": self.camera_id,
                "stream_url": self.stream_url,
                "target_fps": self.target_fps,
                "actual_fps": self.cap.get(cv2.CAP_PROP_FPS),
                "width": int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
                "height": int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT)),
                "codec": self.cap.get(cv2.CAP_PROP_FOURCC),
                "stats": self.stats.copy()
            }
        except Exception as e:
            logger.error(f"获取摄像头 {self.camera_id} 流信息失败: {str(e)}")
            return {
                "status": "error",
                "camera_id": self.camera_id,
                "stream_url": self.stream_url,
                "error": str(e)
            }
            
    async def extract_single_frame(self) -> Optional[np.ndarray]:
        """
        提取单个视频帧
        
        Returns:
            Optional[np.ndarray]: 视频帧，如果失败返回None
        """
        if not self.is_running or not self.cap:
            return None
            
        try:
            ret, frame = self.cap.read()
            if ret and frame is not None:
                self.stats["frames_extracted"] += 1
                self.stats["last_successful_frame"] = time.time()
                return frame
            else:
                self.stats["frames_dropped"] += 1
                return None
        except Exception as e:
            logger.error(f"摄像头 {self.camera_id} 提取单帧失败: {str(e)}")
            self.stats["frames_dropped"] += 1
            return None
            
    def is_healthy(self) -> bool:
        """
        检查流健康状态
        
        Returns:
            bool: 流是否健康
        """
        if not self.is_running or not self.cap or not self.cap.isOpened():
            return False
            
        # 检查最后成功帧的时间
        if self.stats["last_successful_frame"]:
            time_since_last_frame = time.time() - self.stats["last_successful_frame"]
            if time_since_last_frame > settings.STREAM_TIMEOUT:
                return False
                
        return True


@asynccontextmanager
async def create_frame_extractor(stream_url: str, target_fps: int = None, camera_id: str = None):
    """
    创建视频帧提取器的异步上下文管理器
    
    Args:
        stream_url: 视频流URL
        target_fps: 目标帧率
        camera_id: 摄像头ID
        
    Yields:
        VideoFrameExtractor: 视频帧提取器实例
    """
    extractor = VideoFrameExtractor(stream_url, target_fps, camera_id)
    try:
        if await extractor.start():
            yield extractor
        else:
            raise RuntimeError(f"无法启动摄像头 {camera_id} 的视频流")
    finally:
        await extractor.stop()